from typing import Any

from fastapi import APIRouter
from starlette.endpoints import WebSocketEndpoint
from starlette.websockets import WebSocket,WebSocketDisconnect

router = APIRouter()

class Echo(WebSocketEndpoint):
    encoding = "json"
    active_connections = []

    #重写连接方法,  web_socket是对象实例, 表示每一个连接
    async def on_connect(self, web_socket: WebSocket):
        u_type = web_socket.query_params.get("u_type")
        token = web_socket.headers.get("sec-websocket-protocol")
        real_ip = web_socket.headers.get("origin")
        real_host = web_socket.headers.get("host")

    #重写接收方法
    async def on_receive(self, web_socket: WebSocket, msg: Any):
        pass

    #重写断开方法
    async def on_disconnect(self, web_socket: WebSocket, close_code: int):
        pass


router.add_websocket_route('/websockettest', Echo)